import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.patches import Circle
import heapq
from collections import defaultdict


# 节点类
class Node:
    def __init__(self, freq, symbol='', left=None, right=None):
        self.freq = freq
        self.symbol = symbol
        self.left = left
        self.right = right
    # 当创建一个新的`Node` 对象时，这个方法会被调用。
    # `freq`: 这是节点的频率，是一个必需的参数。
    # `symbol`: 这是节点的符号或值，默认为空字符串。
    # `left` 和`right`: 这两个参数分别表示节点的左子节点和右子节点，默认为`None`。

    def __lt__(self, other):
        return self.freq <= other.freq

    # 在这个方法中，我们简单地比较两个节点的 `freq` 属性。如果当前节点的 `freq` 小于等于另一个节点的 `freq`，则返回 `True`，否则返回 `False`。


    # 初始化数据
data = [1, 3, 2, 4, 3, 4, 5]
freq_dict = defaultdict(int)
# defaultdict是Python的collections模块中的一个类，它允许你指定一个默认值，当字典中没有某个键时，会自动创建该键并赋值为指定的默认值。在这个例子中，我们指定了默认值为int（即0），所以当我们试图访问一个不存在的键时，它将自动返回0。

for item in data:
    freq_dict[item] += 1
# req_dict将包含列表中每个元素的出现次数。对于给定的data列表，freq_dict将会是：
# {1: 1, 3: 2, 2: 1, 4: 2, 5: 1}

# 构建哈夫曼树
def build_huffman_tree():
    # 初始化堆：
    heap = [Node(freq, symbol) for symbol, freq in freq_dict.items()]
    # 首先，为每个字符创建一个节点（Node），并将其放入一个列表中。

    heapq.heapify(heap)
    # 然后，使用 heapq.heapify 将这个列表转换成一个堆（heap），以便能够高效地找到和移除最小的元素。
    trees = []

    # 当堆中还有多于一个节点时，继续构建霍夫曼树。
    while len(heap) > 1:

        # 取出频数最小的两个节点
        node1 = heapq.heappop(heap)
        node2 = heapq.heappop(heap)
        # 使用 heapq.heappop 来取出最小元素。
        merged = Node(node1.freq + node2.freq, left=node1, right=node2)
        # 创建一个新的节点，其频率是两个子节点频率的总和，并将这两个子节点作为新节点的左右子树。
        heapq.heappush(heap, merged)
        # 使用 heapq.heappush 将合并后的节点放回堆中
        trees.append(merged)  # 保存每一步合并的树

    trees.append(heap[0])  # 保存最终的树
    return trees


# 绘制函数
def draw(tree, ax, x=0, y=0, depth=0, color='blue'):

    if tree is None:
        return
    # 如果传入的tree是None，则函数直接返回，不执行任何绘制操作。
        # 绘制彩色圈
    circle = Circle((x, y), 8, color=color, fill=True, edgecolor='black')
    ax.add_patch(circle)
    # 使用matplotlib的Circle类在指定坐标(x, y)处绘制一个半径为8的彩色圆圈，表示树的节点。
    # 圆圈被填充为指定的颜色，并且边缘颜色设置为黑色。
    # 然后，将该圆圈添加到轴对象上。
    ax.text(x, y, f'{tree.symbol}:{tree.freq}' if tree.symbol else f'[{tree.freq}]', fontsize=12, ha='center',
            va='center', color='white')
    # 在圆圈的中心位置添加文本。
    # 如果节点有符号（tree.symbol），则显示符号和频率；否则，只显示频率。文本的颜色设置为白色。

    # 使用depth来增加水平间距
    spacing = 20 * depth  # 例如，每个深度级别增加20个单位的间距
    # 根据节点的深度来调整子节点的水平间距。深度越深，间距越小，使得树形结构更加紧凑。

    if tree.left:
        draw(tree.left, ax, x - (50 - spacing), y - 40, depth + 1, 'blue')
        ax.plot([x, x - (50 - spacing)], [y, y - 40], 'k-')
        ax.text(x - (50 - spacing) / 2, (2*y - 40)/2, '0', fontsize=12, ha='center', va='top', color='black')
    # 如果当前节点有左子树，则递归地调用draw函数来绘制左子树。
    # 左子树的根节点位于当前节点的左侧，并稍微向下移动。
    # 然后，使用黑色线条连接当前节点和左子树的根节点。并在线条上添加“0”.

    if tree.right:
        draw(tree.right, ax, x + (50 - spacing), y - 40, depth + 1, 'green')
        ax.plot([x, x + (50 - spacing)], [y, y - 40], 'k-')
        ax.text(x + (50 - spacing) / 2, (2*y - 40)/2, '1', fontsize=12, ha='center', va='top', color='black')
    # 如果当前节点有右子树，则递归地调用draw函数来绘制右子树。
    # 右子树的根节点位于当前节点的右侧，并稍微向下移动。
    # 然后，使用黑色线条连接当前节点和右子树的根节点。并在线条上添加“1”


    # 动画更新函数

# 这个函数用于更新动画的每一帧。

# frame 参数表示当前的帧数，它是一个整数，对应于trees列表中的索引。
# ax 是matplotlib的轴对象，用于绘制图形。
# trees 是一个包含霍夫曼树构建过程中每一步骤的树结构的列表。
def update(frame, ax, trees):
    # 函数首先使用 ax.clear() 清除轴上的所有内容，为绘制新的树做准备。
    ax.clear()
    # 调用 draw 函数来绘制当前帧对应的霍夫曼树
    draw(trees[frame], ax)
    # 设置轴的x和y限制，以确保树在动画中保持适当的大小和位置。
    ax.set_xlim(-100, 100)
    ax.set_ylim(-150, 50)
    # 使用 ax.axis('off') 关闭坐标轴，使图形更美观。
    ax.axis('off')


# 主程序
# 使用 plt.subplots() 创建一个图形（fig）和一个轴（ax）
fig, ax = plt.subplots()
trees = build_huffman_tree()

# 创建动画
ani = FuncAnimation(fig, update, frames=range(len(trees)), fargs=(ax, trees), interval=1000, repeat=False)
# fig 是要在其中创建动画的图形。
# update 是每一帧被调用时执行的函数。
# frames=range(len(trees)) 指定动画的帧数，这里是通过trees列表的长度来确定的。
# fargs=(ax, trees) 指定传递给update函数的额外参数。
# interval=1000 设置动画帧之间的间隔时间（以毫秒为单位），这里是1000毫秒，即1秒。
# repeat=False 指定动画不重复播放。
plt.show()